from vectors import *
from PointEntity import *
from Tkinter import *
from ObjectEntity import *

import numpy

class Floor:
	"""The class which encapsulates the data necessary to display the
	heatmap as elements of a display."""
	name = None
	def __init__(self, id, heatmap, nxlo, nxhi, nylo, nyhi, nzlo, nzhi, \
	             scalemin, scalemax):
		self.id = id
		self.heatmap = heatmap
		if (type(nxlo) != int) or (type(nxhi) != int) or \
		   (type(nylo) != int) or (type(nyhi) != int) or \
		   (type(nzlo) != int) or (type(nzhi) != int):
			raise TypeError("All slice indices must be integers.")
		self.min_scale  = vec2(min(scalemin[0], scalemax[0]), \
		                       min(scalemin[1], scalemax[1]))
		self.max_scale  = vec2(max(scalemin[0], scalemax[0]), \
		                       max(scalemin[1], scalemax[1]))
		# Prevent slicing errors born of negative lower bounds
		if nxlo < 0:
			self.nxlo = 0
		else:
			self.nxlo = nxlo

		if nylo < 0:
			self.nylo = 0
		else:
			self.nylo = nylo

		if nzlo < 0:
			self.nzlo = 0
		else:
			self.nzlo = nzlo

		# Prevent slicing errors born of excessive upper bounds
		if nxhi >= heatmap.n.x:
			self.nxhi = heatmap.n.x
		else:
			self.nxhi = nxhi

		if nyhi >= heatmap.n.y:
			self.nyhi = heatmap.n.y
		else:
			self.nyhi = nyhi

		if nzhi >= heatmap.n.z:
			self.nzhi = heatmap.n.z
		else:
			self.nzhi = nzhi
		
		self.resolution = vec3(self.nxhi - self.nxlo, \
		                       self.nyhi - self.nylo, \
		                       self.nzhi - self.nzlo)
		# this exploits a flaw in vec3s that allows you to reassign elements
		# to non-floats.  Also, self.n is quicker to type than self.resolution.
		# also, this flaw is now a feature.
		self.n = self.resolution
		self.n = self.n.int()
		
		# Compute the volume encompassed by this Floor.
		self.min_corner = heatmap.voxelPos(self.nxlo, \
		                                   self.nylo, \
		                                   self.nzlo) \
		                  - heatmap.voxel_size * .5
		self.max_corner = heatmap.voxelPos(self.nxhi - 1, \
		                                   self.nyhi - 1, \
		                                   self.nzhi - 1) \
		                  + heatmap.voxel_size * .5
		self.size = self.max_corner - self.min_corner
		self.scalesize = self.max_scale - self.min_scale

		self.likelihood = heatmap.likelihood[self.nxlo:self.nxhi, \
		                                       self.nylo:self.nyhi, \
		                                       self.nzlo:self.nzhi]
		
		self.tile_prob = self.likelihood.sum(2)
		# tiles will be the handles of the squares which make up this floor's
		# representation of the heatmap subarray
		temp = [None]*self.n.x
		self.tiles = [temp[:] for i in range(self.n.y)]
	
	def __repr__(self):
		"""How the Floor would be referred to."""
		namestr = ''
		if self.name:
			namestr = " (%s)" % self.name
		return ("Floor %d" % self.id) + namestr

	def tile_prob_update(self):
		"""Gets the most recent probabilities from the heatmap subarray,
		storing them in tile_prob."""
		# This collapses the array along the z dimension.
		self.tile_prob = self.likelihood.sum(2)

	def contains(self, point):
		"""Returns True if this floor contains the given location in space.
		Takes either something vec3-friendly or a PointEntity."""
		if not (is_vec3_friendly(point) or isinstance(point, PointEntity)):
			raise ValueError("Expected the argument of contains to be " + \
			                 "vec3-friendly or a PointEntity, instead " + \
			                 "got %s" % point)
		else:
			# if it's vec3-like, let us access its x, y, z
			if is_vec3_friendly(point):
				p = vec3(point)
			# otherwise, just find the PointEntity's position attribute
			else:
				p = point.p
			if (p.x < self.min_corner.x) or (p.x >= self.max_corner.x) or \
			   (p.y < self.min_corner.y) or (p.y >= self.max_corner.y) or \
			   (p.z < self.min_corner.z) or (p.z >= self.max_corner.z):
				return False
			else:
				return True
	
	def env2canv(self, point):
		"""Given a point, computes the location that such a point would
		occupy in the canvas relative to this floor.  Returns a tuple, not
		a vec2, because of Visuals.initialize_tiles."""
		if not (is_vec3_friendly(point) or is_vec2_friendly(point) or \
		       isinstance(point, PointEntity)):
			raise ValueError("Expected the argument of contains to be " + \
			                 "vec3-friendly or a PointEntity, instead " + \
			                 "got %s" % point)
		# if it's vec3-like, let us access its x, y, z
		if is_vec3_friendly(point):
			p = vec3(point)
		# same goes for vec2-like
		elif is_vec2_friendly(point):
			p = vec2(point)
		# otherwise, just find the PointEntity's position attribute
		else:
			p = point.p
		# TODO: check for correctness.
		canvas_x = self.min_scale.x + (p.x - self.min_corner.x) * \
		           (self.scalesize.x / self.size.x)
		canvas_y = self.max_scale.y - (p.y - self.min_corner.x) * \
		           (self.scalesize.y / self.size.y)
		return (canvas_x, canvas_y)

	def collapseObject(self, object):
		"""An obstacle may overlap partially with a floor.  To avoid drawing the
		bits which are not in the floor, we compute the subobject which should
		be drawn within a floor.
		Returns the min and max corner (in visualization coordinates) of the 
		subobject if there is an intersection with this floor.
		Returns False if there is no intersection with this floor."""
		if not isinstance(object, ObjectEntity):
			raise TypeError("Expected ObjectEntity for argument, instead" + \
			                "got a %s" % type(object))
		# Note that we need the z comparison for intersection, but not for
		# conversion to 2D coordinates.
		if (self.min_corner.x > object.pt2.x) or \
		   (self.max_corner.x < object.pt1.x) or \
		   (self.min_corner.y > object.pt2.y) or \
		   (self.max_corner.y < object.pt1.y) or \
		   (self.min_corner.z > object.pt2.z) or \
		   (self.max_corner.z < object.pt1.z):
			return False
		else:
			return (self.env2canv((max(self.min_corner.x, object.pt1.x), \
			                       max(self.min_corner.y, object.pt1.y))), \
			        self.env2canv((min(self.max_corner.x, object.pt2.x), \
			                       min(self.max_corner.y, object.pt2.y))))





class Visuals:
	"""This class provides the interface with Tkinter.  It links directly to
	the Simulator class, so that it may access the positions of all objects
	and specifications in the environment, as well as the probabilities in 
	the heatmap."""
	def __init__(self, simulator):
		self.sim = simulator
		self.debug_mode = simulator.debug_mode
		self.heatmap = simulator.heatmap
		self.mapdata = simulator.mapdata
		self.floors  = simulator.floors
		self.tk = Tk()
		self.tk.geometry(simulator.window_geometry)
		self.canvas = Canvas(self.tk,
		                     width=self.sim.canvas_dim.x,
		                     height=self.sim.canvas_dim.y)
		self.c = self.canvas
		self.tk.title(self.sim.map_file)
		self.initialize_tiles()
		# The detector and emitter handles must be in ordered lists, because
		# it is not guaranteed that they will occupy exactly one floor.
		# The lists are organized list[floor][somethingID].
		temp = [None] * len(self.sim.obstacles)
		self.obstacle_handles = [temp[:] for i in range(len(self.floors))]
		temp = [None] * len(self.sim.detectors)
		self.detector_handles = [temp[:] for i in range(len(self.floors))]
		temp = [None] * len(self.sim.emitters)
		self.emitter_handles = [temp[:] for i in range(len(self.floors))]
		self.initialize_obstacles()
		self.redraw_detectors()
		self.redraw_emitters()
		
		self.canvas.pack()
	
	def msg(self, message):
		"""A message that will only be printed if debug_mode is set to True."""
		if self.debug_mode: print message
	
	def update(self, subtractMean=False):
		"""Retrieves the most recent set of data and redraws the 
		display accordingly."""
		self.update_tiles(subtractMean)
		self.redraw_detectors()
		self.redraw_emitters()
		self.tk.update()
		self.msg("Visualization update complete. Time = %s" % self.sim.time)

	def initialize_tiles(self):
		"""Figures out where onscreen the heatmap tiles should be placed and
		draws them with their current likelihood estimates."""
		vs = self.heatmap.voxel_size
		min_c = self.heatmap.min_corner
		max_prob = self.heatmap.likelihood.max()
		for f in self.floors:
			floor = self.floors[f]
			for i in range(floor.n.x):
				for j in range(floor.n.y):
					x1 = min_c.x + vs.x *  i
					x2 = min_c.x + vs.x * (i+1)
					y1 = min_c.y + vs.y *  j
					y2 = min_c.y + vs.y * (j+1)
					c1 = floor.env2canv((x1, y1))
					c2 = floor.env2canv((x2, y2))
					floor.tiles[i][j] = \
						self.c.create_rectangle(c1, c2, \
							fill=prob2red(floor.tile_prob[i][j], \
							              ),#normalizeTo=max_prob), \
							outline="")

	def update_tiles(self, subtractMean=False):
		"""Collapses the array of probabilities into color intensities for
		each floor and redraws all floors."""
		max_tile = 0.0
		for f in self.floors:
			floor = self.floors[f]
			floor.tile_prob_update()
			max_tile = max(max_tile, floor.tile_prob.max())
		
		for f in self.floors:
			floor = self.floors[f]
			floor.tile_prob_update()
			if subtractMean:
				temp = floor.tile_prob - floor.tile_prob.mean()
				max_tile = temp.max()
			else:
				temp = floor.tile_prob
			for i in range(floor.n.x):
				for j in range(floor.n.y):
					self.c.itemconfig(floor.tiles[i][j], \
					                  fill=prob2red(temp[i][j], \
					                                normalizeTo=max_tile))

	def redraw_detectors(self, \
	                   sidelength=4, \
	                   default_color="#008800", \
	                   hot_color="#FFFF00",
	                   direction=True):
		"""Figures out where onscreen the detectors should be placed and
		draws them.  Because detectors may move from floor to floor, or be
		on more than one 'Floor' at a time, this method erases and redraws
		the symbols rather than attempting to move them.
		TODO: figure out if this is actually a decent implementation"""
		floors = self.floors
		detectors = self.sim.detectors
		# Find the first group_size detectors in the ordered list so we know
		# which ones to color brighter, but only if we think there's a
		# radioactive anomaly somewhere in the field.
		if self.sim.source_detected:
			self.msg(self.sim.hottest_groups)
			hot_group = self.sim.hottest_groups[4]
		else:
			hot_group = []
		self.erase_detectors()
		for i in range(len(floors)):
			for j in range(len(detectors)):
				if floors[i].contains(detectors[j]):
					if j in hot_group:
						color = hot_color
					else:
						color = default_color
					self.detector_handles[i][j] = \
						draw_square(self.c, \
						            floors[i].env2canv(detectors[j].p), \
						            sidelength, \
						            color)
		self.c.update()
		
	def erase_detectors(self):
		"""Intended for use within draw_detectors.  This method removes all
		detectors from the viewscreen which are present on any floor."""
		for i in range(len(self.floors)):
			for j in range(len(self.sim.detectors)):
				if self.detector_handles[i][j]:
					self.c.delete(self.detector_handles[i][j])
					self.detector_handles[i][j] = None
		self.c.update()

	def redraw_emitters(self, radius=4, color="#FFFFFF"):
		"""Figures out where onscreen the emitters should be placed and
		draws them.  Because emitters may move from floor to floor, or be
		on more than one 'Floor' at a time, this method erases and redraws
		the symbols rather than attempting to move them.
		TODO: figure out if this is actually a decent implementation"""
		floors = self.floors
		emitters = self.sim.emitters
		self.erase_emitters()
		for i in range(len(floors)):
			for j in range(len(emitters)):
				if floors[i].contains(emitters[j]):
					self.emitter_handles[i][j] = \
						draw_circle(self.c, \
						            floors[i].env2canv(emitters[j].p), \
						            radius, \
						            color)
		self.c.update()
		
	def erase_emitters(self):
		"""Intended for use within draw_emitters.  This method removes all
		emitters from the viewscreen which are present on any floor."""
		for i in range(len(self.floors)):
			for j in range(len(self.sim.emitters)):
				if self.emitter_handles[i][j]:
					self.c.delete(self.emitter_handles[i][j])
					self.emitter_handles[i][j] = None
		self.c.update()

	def initialize_obstacles(self, color="#0040CC"):
		"""Figures out where onscreen the obstacles should be placed and
		draws them.
		TODO: Implement."""
		floors = self.floors
		obstacles = self.sim.obstacles
		for i in range(len(floors)):
			for j in range(len(obstacles)):
				# Box is just a 2-tuple of 2-tuples,
				# ((pt1.x, pt1.y), (pt2.x, pt2.y))
				box = floors[i].collapseObject(obstacles[j])
				if box:
					self.obstacle_handles[i][j] = \
						self.c.create_rectangle(box[0][0], box[0][1], \
						                        box[1][0], box[1][1], \
						                        fill = "", \
						                        outline = color)


def draw_square(canvas, center, side_length, color):
	"""Draws a square on the canvas, centered at center."""
	if not is_vec2_friendly(center):
		return TypeError("tile_size and center must be vec2-friendly.")
	# (x1, y1) is the position of the lower left corner of the square
	x1 = center[0] - side_length/2.0
	y1 = center[1] - side_length/2.0
	# (x2, y2) is the position of the upper right corner of the square
	x2 = x1 + side_length
	y2 = y1 + side_length
	return canvas.create_rectangle(x1, y1, \
	                               x2, y2, \
	                               fill = color, \
	                               outline = color)

def draw_circle(canvas, center, radius, color):
	"""Draws a circle on the canvas, centered at center, with the specified
	radius, of the specified color"""
	# (x1, y1) is the position of the lower left corner of the circle
	x1 = center[0] - radius
	y1 = center[1] - radius
	# (x2, y2) is the position of the upper right corner of the circle
	x2 = x1 + 2*radius
	y2 = y1 + 2*radius
	return canvas.create_oval(x1, y1, x2, y2, \
	                          fill = color, \
	                          outline = color)

def prob2red(prob, normalizeTo=1.0):
	"""Given a probability (zero to one) translates it into a hex color code
	string (some intensity of red) and returns it."""
	# normalizeTo is in case of very low probabilities, and some desire to 
	# make the background not black.
	if prob <= 0:
		return "#000000"
	if normalizeTo <= 0.0:
		normalizeTo = 1.0
	red = "%x" % int(prob*255/normalizeTo)
	if len(red) == 1:
		red = "0" + red
	# if normalization pushes the red proportion over 255, make it 255.
	elif len(red) > 2:
		red = "ff"
	return "#%s0000" % red